# -*-coding:utf-8-*-

import numpy as np
import os
from time import sleep


# KNN分类核心方法 newInput测试集的一行数据 dataSet训练集 labels标签 k
def kNNClassify(newInput, dataSet, labels, k):
    numSample = len(dataSet)
    print(dataSet.shape[0])
    print(numSample)
    privVec = np.tile(newInput, (numSample, 1)) - dataSet
    diff = privVec ** 2
    print(type(diff[0][0]))
    diffTmp = np.sum(diff, axis=1)
    diffResult = diffTmp ** 0.5
    arg = np.argsort(diffResult)
    classDict = {}
    for i in range(k):
        vot = labels[arg[i]]
        classDict[vot] = classDict.get(vot, 0) + 1
    print(classDict)

    maxNum = 0
    maxIndex = 0
    for key, value in classDict.items():
        if value > maxNum:
            maxNum = value
            maxIndex = key
    print(maxIndex)
    return maxIndex


# 将图片转换为向量
def img2vector(filename):
    rows = 32
    cols = 32
    file = open(filename)
    x = np.zeros((1, rows * cols))
    for row in range(rows):
        line = file.readline()
        for col in range(cols):
            x[0][row * 32 + col] = int(line[col])
    return x


# 加载数据集
def loadDataSet():
    files = os.listdir("./train")
    numSample = len(files)
    trainX = np.zeros((numSample, 32 * 32))
    trainY = []
    for i in range(numSample):
        filename = files[i]
        trainX[i, :] = img2vector("./train/" + filename)
        trainY.append(int(filename.split("_")[0]))
    testFiles = os.listdir("./test")
    testSample = len(testFiles)
    testX = np.zeros((testSample, 32 * 32))
    testY = []
    for i in range(testSample):
        filename = testFiles[i]
        testX[i, :] = img2vector("./test/" + filename)
        testY.append(int(filename.split("_")[0]))
    return trainX, trainY, testX, testY


trainX, trainY, testX, testY = loadDataSet()
for i in range(len(testX)):
    sleep(1)
    print("预测值：", kNNClassify(testX[i], trainX, trainY, 3), "   真实值:", testY[i])
